import numpy as np
from numba import jit, prange, float64, types
from typing import Tuple

# 基础物理计算函数
@jit(float64(float64), nopython=True, parallel=False)
def calculate_friction_factor(Re: float) -> float:
    """计算壁面摩擦因子"""
    return 0.184 * Re**(-0.2)

@jit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:]), nopython=True, parallel=True)
def calculate_mass_flux_array(alpha: np.ndarray, rho_g: np.ndarray, rho_l: np.ndarray, 
                            v_g: np.ndarray, v_l: np.ndarray) -> np.ndarray:
    """并行计算质量流量"""
    n = len(alpha)
    result = np.zeros(2*n)  # [G_g, G_l]交错存储
    
    for i in prange(n):
        result[2*i] = alpha[i] * rho_g[i] * v_g[i]      # G_g
        result[2*i+1] = (1 - alpha[i]) * rho_l[i] * v_l[i]  # G_l
    
    return result

@jit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:]), nopython=True, parallel=True)
def calculate_momentum_flux_array(alpha: np.ndarray, rho_g: np.ndarray, rho_l: np.ndarray,
                                v_g: np.ndarray, v_l: np.ndarray) -> np.ndarray:
    """并行计算动量流量"""
    n = len(alpha)
    result = np.zeros(2*n)  # [M_g, M_l]交错存储
    
    for i in prange(n):
        result[2*i] = alpha[i] * rho_g[i] * v_g[i] * v_g[i]      # M_g
        result[2*i+1] = (1 - alpha[i]) * rho_l[i] * v_l[i] * v_l[i]  # M_l
    
    return result

@jit(float64[:](float64[:], float64[:], float64[:], float64[:]), nopython=True, parallel=True)
def calculate_wall_friction_array(rho: np.ndarray, v: np.ndarray, D_h: np.ndarray, 
                                mu: np.ndarray) -> np.ndarray:
    """并行计算壁面摩擦力"""
    n = len(rho)
    F_w = np.zeros(n)
    
    for i in prange(n):
        Re = abs(rho[i] * v[i] * D_h[i] / mu[i])
        f = calculate_friction_factor(Re)
        F_w[i] = -0.5 * f * rho[i] * v[i] * abs(v[i]) / D_h[i]
    
    return F_w

@jit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64[:]), 
     nopython=True, parallel=True)
def calculate_interfacial_friction_array(alpha: np.ndarray, rho_g: np.ndarray, rho_l: np.ndarray,
                                       v_g: np.ndarray, v_l: np.ndarray, D_h: np.ndarray) -> np.ndarray:
    """并行计算界面摩擦力"""
    n = len(alpha)
    F_i = np.zeros(n)
    C_i = 0.5  # 界面摩擦系数
    
    for i in prange(n):
        v_r = v_g[i] - v_l[i]
        F_i[i] = -C_i * alpha[i] * (1 - alpha[i]) * rho_l[i] * v_r * abs(v_r) / D_h[i]
    
    return F_i

@jit(float64[:](float64[:], float64[:], float64[:]), nopython=True, parallel=True)
def calculate_lateral_flow_resistance_array(rho: np.ndarray, v: np.ndarray, s: np.ndarray) -> np.ndarray:
    """并行计算横向流动阻力"""
    n = len(rho)
    R = np.zeros(n)
    K_ij = 0.5  # 间隙横向流阻
    
    for i in prange(n):
        R[i] = -0.5 * K_ij * rho[i] * v[i] * abs(v[i]) / s[i]
    
    return R

@jit(float64[:](float64, float64[:], float64[:], float64[:], float64[:]), nopython=True, parallel=True)
def calculate_turbulent_mixing_array(beta: float, s: np.ndarray, rho_1: np.ndarray, 
                                   rho_2: np.ndarray, v_1: np.ndarray) -> np.ndarray:
    """并行计算湍流交混项"""
    n = len(s)
    w_t = np.zeros(n)
    
    for i in prange(n):
        rho_avg = 0.5 * (rho_1[i] + rho_2[i])
        v_avg = 0.5 * abs(v_1[i])
        w_t[i] = beta * s[i] * rho_avg * v_avg
    
    return w_t

@jit(types.Tuple((float64[:], float64[:]))(float64[:], float64[:], float64[:], float64[:], float64[:], 
                                    float64[:], float64, float64), nopython=True, parallel=True)
def solve_momentum_equation_parallel(v_old: np.ndarray, p: np.ndarray, rho: np.ndarray, 
                                   F_w: np.ndarray, F_i: np.ndarray, alpha: np.ndarray,
                                   dt: float, dz: float) -> Tuple[np.ndarray, np.ndarray]:
    """并行求解动量方程"""
    n = len(v_old)
    v_new = np.zeros_like(v_old)
    residual = np.zeros_like(v_old)
    
    for i in prange(1, n-1):
        # 计算压力梯度
        dp_dz = (p[i+1] - p[i-1]) / (2 * dz)
        
        # 计算速度
        v_new[i] = v_old[i] + dt * (
            -dp_dz / rho[i] +
            F_w[i] / (rho[i] * alpha[i]) +
            F_i[i] / (rho[i] * alpha[i]) -
            9.81  # 重力项
        )
        
        # 计算残差
        residual[i] = abs(v_new[i] - v_old[i])
    
    # 边界条件处理
    v_new[0] = v_old[0]
    v_new[-1] = v_old[-1]
    residual[0] = 0.0
    residual[-1] = 0.0
    
    return v_new, residual

@jit(float64[:](float64[:], float64[:], float64[:], float64[:], float64[:], float64,
                float64), nopython=True, parallel=True)
def solve_continuity_equation_parallel(alpha_old: np.ndarray, G_g: np.ndarray, G_l: np.ndarray,
                                     rho_g: np.ndarray, rho_l: np.ndarray, dt: float,
                                     area: float) -> np.ndarray:
    """并行求解连续性方程"""
    n = len(alpha_old)
    alpha_new = np.zeros_like(alpha_old)
    
    for i in prange(1, n-1):
        # 计算空泡份额
        alpha_flux = (G_g[i] / rho_g[i] - G_l[i] / rho_l[i])
        alpha_new[i] = alpha_old[i] + dt * (-alpha_flux / area)
        
        # 限制空泡份额在物理范围内
        alpha_new[i] = max(0.0, min(1.0, alpha_new[i]))
    
    # 边界条件处理
    alpha_new[0] = alpha_old[0]
    alpha_new[-1] = alpha_old[-1]
    
    return alpha_new 